From cb1e615715a5d52960105a3972a1b7822b901688 Mon Sep 17 00:00:00 2001
From: Karol Kuczmarski <karol.kuczmarski@gmail.com>
Date: Sat, 13 May 2017 16:51:40 +0100
Subject: [PATCH] Fix build script envvars for `cargo test/doc`

---
 src/cargo/ops/cargo_rustc/mod.rs | 16 +++++++++---
 tests/build-script.rs            | 44 ++++++++++++++++++++++++++------
 2 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/src/cargo/ops/cargo_rustc/mod.rs b/src/cargo/ops/cargo_rustc/mod.rs
index bf46397c7..470f700f5 100644
--- a/src/cargo/ops/cargo_rustc/mod.rs
+++ b/src/cargo/ops/cargo_rustc/mod.rs
@@ -195,6 +195,10 @@ pub fn compile_targets<'a, 'cfg: 'a>(ws: &Workspace<'cfg>,
             .or_insert_with(HashSet::new)
             .extend(output.cfgs.iter().cloned());
 
+        cx.compilation.extra_env.entry(pkg.clone())
+            .or_insert_with(Vec::new)
+            .extend(output.env.iter().cloned());
+
         for dir in output.library_paths.iter() {
             cx.compilation.native_dirs.insert(dir.clone());
         }
@@ -265,6 +269,7 @@ fn rustc(cx: &mut Context, unit: &Unit, exec: Arc<Executor>) -> CargoResult<Work
 
     let filenames = cx.target_filenames(unit)?;
     let root = cx.out_dir(unit);
+    let kind = unit.kind;
 
     // Prepare the native lib state (extra -L and -l flags)
     let build_state = cx.build_state.clone();
@@ -313,7 +318,7 @@ fn rustc(cx: &mut Context, unit: &Unit, exec: Arc<Executor>) -> CargoResult<Work
                                  pass_l_flag, &current_id)?;
             add_plugin_deps(&mut rustc, &build_state, &build_deps,
                                  &root_output)?;
-            add_custom_env(&mut rustc, &build_state, &build_deps, &current_id)?;
+            add_custom_env(&mut rustc, &build_state, &current_id, kind)?;
         }
 
         // FIXME(rust-lang/rust#18913): we probably shouldn't have to do
@@ -428,9 +433,9 @@ fn rustc(cx: &mut Context, unit: &Unit, exec: Arc<Executor>) -> CargoResult<Work
     // been put there by one of the `build_scripts`) to the command provided.
     fn add_custom_env(rustc: &mut ProcessBuilder,
                       build_state: &BuildMap,
-                      _: &BuildScripts,
-                      current_id: &PackageId) -> CargoResult<()> {
-        let key = (current_id.clone(), Kind::Host);
+                      current_id: &PackageId,
+                      kind: Kind) -> CargoResult<()> {
+        let key = (current_id.clone(), kind);
         if let Some(output) = build_state.get(&key) {
             for &(ref name, ref value) in output.env.iter() {
                 rustc.env(name, value);
@@ -617,6 +622,9 @@ fn rustdoc(cx: &mut Context, unit: &Unit) -> CargoResult<Work> {
             for cfg in output.cfgs.iter() {
                 rustdoc.arg("--cfg").arg(cfg);
             }
+            for &(ref name, ref value) in output.env.iter() {
+                rustdoc.env(name, value);
+            }
         }
         state.running(&rustdoc);
         rustdoc.exec().chain_error(|| {
diff --git a/tests/build-script.rs b/tests/build-script.rs
index 31dbd7b46..7356af284 100644
--- a/tests/build-script.rs
+++ b/tests/build-script.rs
@@ -1577,10 +1577,10 @@ fn cfg_override_doc() {
 
 #[test]
 fn env_build() {
-    let build = project("builder")
+    let p = project("foo")
         .file("Cargo.toml", r#"
             [package]
-            name = "builder"
+            name = "foo"
             version = "0.0.1"
             authors = []
             build = "build.rs"
@@ -1596,9 +1596,9 @@ fn env_build() {
                 println!("cargo:rustc-env=FOO=foo");
             }
         "#);
-    assert_that(build.cargo_process("build").arg("-v"),
+    assert_that(p.cargo_process("build").arg("-v"),
                 execs().with_status(0));
-    assert_that(build.cargo("run").arg("-v"),
+    assert_that(p.cargo("run").arg("-v"),
                 execs().with_status(0).with_stdout("foo\n"));
 }
 
@@ -1633,14 +1633,14 @@ fn env_test() {
 [COMPILING] foo v0.0.1 ({dir})
 [RUNNING] [..] build.rs [..]
 [RUNNING] `[..][/]build-script-build`
-[RUNNING] [..] --cfg foo[..]
-[RUNNING] [..] --cfg foo[..]
-[RUNNING] [..] --cfg foo[..]
+[RUNNING] [..] --crate-name foo[..]
+[RUNNING] [..] --crate-name foo[..]
+[RUNNING] [..] --crate-name test[..]
 [FINISHED] dev [unoptimized + debuginfo] target(s) in [..]
 [RUNNING] `[..][/]foo-[..][EXE]`
 [RUNNING] `[..][/]test-[..][EXE]`
 [DOCTEST] foo
-[RUNNING] [..] --cfg foo[..]", dir = p.url()))
+[RUNNING] [..] --crate-name foo[..]", dir = p.url()))
                        .with_stdout("
 running 0 tests
 
@@ -1652,9 +1652,37 @@ test test_foo ... ok
 
 test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured
 
+
+running 0 tests
+
+test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured
+
 "));
 }
 
+#[test]
+fn env_doc() {
+    let p = project("foo")
+        .file("Cargo.toml", r#"
+            [package]
+            name = "foo"
+            version = "0.0.1"
+            authors = []
+            build = "build.rs"
+        "#)
+        .file("src/main.rs", r#"
+            const FOO: &'static str = env!("FOO");
+            fn main() {}
+        "#)
+        .file("build.rs", r#"
+            fn main() {
+                println!("cargo:rustc-env=FOO=foo");
+            }
+        "#);
+    assert_that(p.cargo_process("doc").arg("-v"),
+                execs().with_status(0));
+}
+
 #[test]
 fn flags_go_into_tests() {
     let p = project("foo")
-- 
2.30.2